from config import ALL_MAPPING


def test_nn_Sequential():
    xxx = r"""nn.Sequential(
        nn.Conv2d(1, 20, 5),
        nn.ReLU(),
        nn.Conv2d(20, 64, 5),
        nn.ReLU()
    )"""

    yyy = r"""nn.Sequential(OrderedDict([
        ('conv1', nn.Conv2d(1, 20, 5)),
        ('relu1', nn.ReLU()),
        ('conv2', nn.Conv2d(20, 64, 5)),
        ('relu2', nn.ReLU())
    ]))"""

    zzz = "nn.Sequential(*layers)"
    aaa = "nn.Sequential(layers)"

    print(ALL_MAPPING['nn.Sequential'].convert(xxx))
    print(ALL_MAPPING['nn.Sequential'].convert(yyy))
    print(ALL_MAPPING['nn.Sequential'].convert(zzz))
    print(ALL_MAPPING['nn.Sequential'].convert(aaa))


def test_dot_view():
    expr = "aaa.bbb.ccc.out.view(out.size(0), -1)"
    print(ALL_MAPPING['.view'].convert(expr))


def test_dot_size():
    expr = "out.size()"
    print(ALL_MAPPING['.size'].convert(expr))

    expr = "out.size(111)"
    print(ALL_MAPPING['.size'].convert(expr))

    expr = "out.size(xxx)"
    print(ALL_MAPPING['.size'].convert(expr))


def test_multi_call():
    call_name = "nn.functional.adaptive_avg_pool2d(x, 1)"
    args_str = "(x.shape[0], -1)"
    print(ALL_MAPPING['.reshape'].convert(call_name, args_str))


# test_nn_Sequential()
# test_dot_size()
test_multi_call()

# 找到torchvision/models中包含'nn.Conv2d('的语句
# find /Users/xuxu/miniconda3/envs/wdx_py37/lib/python3.7/site-packages/torchvision/models -name '*.py' | xargs cat | grep -A 3 'nn.Conv2d(' > all.log
